import pandas as pd

from nlp_tools.corpus import BaseCorpusLoader

class DataFoundClassify(BaseCorpusLoader):
    '''
    第五届“达观杯” 基于大规模预训练模型的风险事件标签识别
    https://www.datafountain.cn/competitions/512

    数据格式： label   text
              0       xxxxx
              1       xxxxx
    '''

    def __init__(self, data_path, shuffle: bool = True, split_train_test=False, split_test_rate=0.1, max_length=128):
        super().__init__(data_path, shuffle, split_train_test, split_test_rate, max_length)

    def get_x_and_y(self, df):
        x = df['text']
        y = df['label']
        y = [str(i) for i in y]

        return x,y


if __name__ == '__main__':
    data_path = r'/home/fanfanfeng/PycharmProjects/pythonProject/torch_study/sentiment.csv'

    data = DataFoundClassify(data_path,split_test_rate=True,max_length=32).load_data()
    print(data)